package mcp

import (
	"encoding/json"
	"net/http"
	"github.com/gorilla/websocket"

	"project-info/src/lib")

// MCP协议消息类型
type MCPMessage struct {
	JSONRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id,omitempty"`
	Method  string      `json:"method,omitempty"`
	Params  interface{} `json:"params,omitempty"`
	Result  interface{} `json:"result,omitempty"`
	Error   *MCPError   `json:"error,omitempty"`
}

type MCPError struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

// MCP工具定义（用于协议传输）
type MCPTool struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description"`
	InputSchema map[string]interface{} `json:"inputSchema"`
}

// MCP资源定义
type MCPResource struct {
	URI         string `json:"uri"`
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	MimeType    string `json:"mimeType,omitempty"`
}

// MCP服务器信息
type MCPServerInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// MCP服务器能力
type MCPCapabilities struct {
	Tools     *MCPToolsCapability     `json:"tools,omitempty"`
	Resources *MCPResourcesCapability `json:"resources,omitempty"`
	Prompts   *MCPPromptsCapability   `json:"prompts,omitempty"`
}

type MCPToolsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}

type MCPResourcesCapability struct {
	Subscribe   bool `json:"subscribe,omitempty"`
	ListChanged bool `json:"listChanged,omitempty"`
}

type MCPPromptsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}

// WebSocket升级器
var mcpUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	Subprotocols: []string{"mcp"},
}

// MCP WebSocket处理器
func MCPWebSocketHandler(c *lib.GinContext) error {
	ws, err := mcpUpgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		c.Logger().Error("MCP WebSocket升级失败: " + err.Error())
		return err
	}
	defer ws.Close()

	c.Logger().Info("新的MCP客户端连接")

	// 创建内部客户端用于处理请求
	internalClient := NewMCPInternalClient(c)

	// 处理MCP消息
	for {
		_, messageBytes, err := ws.ReadMessage()
		if err != nil {
			c.Logger().Error("读取MCP消息失败: " + err.Error())
			break
		}

		var msg MCPMessage
		if err := json.Unmarshal(messageBytes, &msg); err != nil {
			c.Logger().Error("解析MCP消息失败: " + err.Error())
			continue
		}

		// 处理MCP方法
		response := handleMCPMethod(c, internalClient, &msg)
		responseBytes, _ := json.Marshal(response)

		if err := ws.WriteMessage(websocket.TextMessage, responseBytes); err != nil {
			c.Logger().Error("发送MCP响应失败: " + err.Error())
			break
		}
	}

	return nil
}

// 处理MCP方法
func handleMCPMethod(c *lib.GinContext, client *MCPInternalClient, msg *MCPMessage) *MCPMessage {
	switch msg.Method {
	case "initialize":
		return handleInitialize(msg)
	case "tools/list":
		return handleToolsList(client, msg)
	case "tools/call":
		return handleToolsCall(client, msg)
	case "resources/list":
		return handleResourcesList(msg)
	case "resources/read":
		return handleResourcesRead(c, msg)
	case "prompts/list":
		return handlePromptsList(msg)
	case "prompts/get":
		return handlePromptsGet(c, msg)
	default:
		return &MCPMessage{
			JSONRPC: "2.0",
			ID:      msg.ID,
			Error: &MCPError{
				Code:    -32601,
				Message: "Method not found",
			},
		}
	}
}

// 初始化处理
func handleInitialize(msg *MCPMessage) *MCPMessage {
	return &MCPMessage{
		JSONRPC: "2.0",
		ID:      msg.ID,
		Result: map[string]interface{}{
			"protocolVersion": "2024-11-05",
			"serverInfo": MCPServerInfo{
				Name:    "project-info-mcp-server",
				Version: "1.0.0",
			},
			"capabilities": MCPCapabilities{
				Tools: &MCPToolsCapability{
					ListChanged: true,
				},
				Resources: &MCPResourcesCapability{
					Subscribe:   false,
					ListChanged: true,
				},
				Prompts: &MCPPromptsCapability{
					ListChanged: true,
				},
			},
		},
	}
}

// 工具列表处理（使用统一的工具管理器）
func handleToolsList(client *MCPInternalClient, msg *MCPMessage) *MCPMessage {
	tools := client.GetAvailableTools()
	return &MCPMessage{
		JSONRPC: "2.0",
		ID:      msg.ID,
		Result: map[string]interface{}{
			"tools": tools,
		},
	}
}

// 工具调用处理（使用统一的工具管理器）
func handleToolsCall(client *MCPInternalClient, msg *MCPMessage) *MCPMessage {
	params, ok := msg.Params.(map[string]interface{})
	if !ok {
		return createMCPErrorResponse(msg.ID, -32602, "Invalid params")
	}

	toolName, ok := params["name"].(string)
	if !ok {
		return createMCPErrorResponse(msg.ID, -32602, "Missing tool name")
	}

	arguments, _ := params["arguments"].(map[string]interface{})
	if arguments == nil {
		arguments = make(map[string]interface{})
	}

	// 使用统一的工具调用
	toolResult := client.CallTool(toolName, arguments)

	if !toolResult.Success {
		return createMCPErrorResponse(msg.ID, -32603, toolResult.Error)
	}

	return &MCPMessage{
		JSONRPC: "2.0",
		ID:      msg.ID,
		Result: map[string]interface{}{
			"content": []map[string]interface{}{
				{
					"type": "text",
					"text": toolResult.Content,
				},
			},
			"isError": false,
		},
	}
}

// 创建MCP错误响应
func createMCPErrorResponse(id interface{}, code int, message string) *MCPMessage {
	return &MCPMessage{
		JSONRPC: "2.0",
		ID:      id,
		Error: &MCPError{
			Code:    code,
			Message: message,
		},
	}
}
